import catboost as cb
from sklearn.metrics import mean_absolute_error,  make_scorer
import pandas as pd
import numpy as np


def main(args):
    cat = {
        'model': 'category',
        'brand': 'category',
        'bodyType': 'category',
        'fuelType': 'category',
        'gearbox': 'category',
        'notRepairedDamage': 'category',
    }

    df = pd.read_csv('./user_data/df_s.csv', sep=' ', dtype=cat)
    train_X = df[df.train == 1].drop(['price', 'SaleID', 'regionCode'], axis=1)
    train_y = df[df.train == 1]['price']
    train_y_ln = np.log1p(train_y)

    model = cb.CatBoostRegressor(**args)
    #######用于交叉验证时将以下代码删去#######
    test_X = df[df.train == 0].drop(['price', 'SaleID', 'regionCode'], axis=1)
    model.fit(X=train_X,
              y=train_y_ln,
              cat_features=[i for i in cat.keys() if i in train_X.columns],
              early_stopping_rounds=300,
              verbose=10000)
    p = model.predict(test_X)
    submit = pd.DataFrame()
    submit['SaleID'] = df[df.train == 0].SaleID
    submit['price'] = np.expm1(p)
    submit.to_csv('./prediction_result/submit_cgb.csv', index=0)
    print('end for predict')
    return
    #######用于交叉验证时将以上代码删去#######
    from sklearn.model_selection import cross_val_score
    n = 1

    def maee(y_true, y_pred):
        nonlocal n
        loss = mean_absolute_error(np.expm1(y_true), np.expm1(y_pred))
        print(f'mae in fold {n}:{loss}')
        n += 1
        return loss
    fit_params = {
        'cat_features': [i for i in cat.keys() if i in train_X.columns],
        'early_stopping_rounds': 300,
        'verbose': 500
    }
    mae = cross_val_score(model, X=train_X, y=train_y_ln, verbose=0,
                          cv=5, scoring=make_scorer(maee), fit_params=fit_params)
    print(mae)
    print(f'with mean mae:{np.mean(mae)}')


if __name__ == '__main__':
    params = {
        'n_estimators': 1000000,
        'loss_function': 'MAE',
        'eval_metric': 'MAE',
        'learning_rate': 0.02,# 设置为0.02会导致训练时间极——————————————长，建议设高一点或者找gpu加速
        'depth': 6,
        # 'use_best_model': True,
        # 'task_type':'GPU'
        'subsample': 0.6,
        'bootstrap_type': 'Bernoulli',
        'reg_lambda': 3,
        'one_hot_max_size': 2,
    }
    main(params)
